#include <stdio.h>
#include <string.h>
#include <time.h>
#include "api.h"
#include "rand.h"
#include "bin-lwe.h"
#include "ntt.h"
#include "ntt-1024.h"
#include "fips202.h"

#ifdef LINUX
#include <x86intrin.h>
#define cpucycles() _rdtsc()
#endif

#ifdef WIN
#include <intrin.h>
#define cpucycles() __rdtsc()
#endif

#define NTESTS 100000

unsigned long long sum, average;

static void print_uint64(unsigned long long num)
{
	if (num >= 10)
		print_uint64(num / 10);
	printf("%d", (int)(num % 10));
}

int print_cpucycles(unsigned long long *t)  // 输出cpu周期
{
	int i;
	sum = 0;
	for (i = 0; i < NTESTS - 1; i++)
	{
		t[i] = t[i + 1] - t[i];
		sum += t[i];
	}
	average = sum / (NTESTS - 1);
	printf(": ");
	print_uint64(average);
	printf(" cpucycles\n");

	// printf("time     :%f us\n",(double)average/2700);

	return 0;
}

// test pke
int test_pke_cpucycles()
{
	printf("\n=============================test_pke_cpucycles==================================\n");
	unsigned long long t[NTESTS];
	unsigned char pk[CRYPTO_PUBLICKEYBYTES];
	unsigned char sk[CRYPTO_SECRETKEYBYTES];
	unsigned char k1[CRYPTO_BYTES], k2[CRYPTO_BYTES], c[CRYPTO_CIPHERTEXTBYTES];
	int i;
	unsigned long long mlen = CRYPTO_BYTES, clen = CRYPTO_CIPHERTEXTBYTES;

	unsigned char seed[SEED_LEN];
	random_bytes(seed,SEED_LEN);
	for (i = 0; i < NTESTS; i++)
	{
		t[i] = cpucycles();
		crypto_encrypt_keypair(pk, sk);
	}
	printf("pke kg     ");
	print_cpucycles(t);

	crypto_encrypt_keypair(pk, sk);
	random_bytes(k1, CRYPTO_BYTES);
    for(i=0;i<NTESTS;i++)
	{
		t[i]=cpucycles();
		crypto_encrypt(c,&clen,k1,mlen,pk,seed);
	}
	printf("encryption ");
	print_cpucycles(t);

	/* FOR POLAR */
	// test polar_enc
	polar_enc_cpucycles();
	/* END FOR POLAR */

	crypto_encrypt_keypair(pk, sk);
	random_bytes(k1, CRYPTO_BYTES);
	crypto_encrypt(c,&clen,k1,mlen,pk, seed);
	for (i = 0; i < NTESTS; i++)
	{
		t[i] = cpucycles();
		crypto_encrypt_open(k2, &mlen, c, clen, sk);
	}
	printf("decryption ");
	print_cpucycles(t);

	/* FOR POLAR */
	// test polar_dec
	polar_dec_cpucycles();
	/* END FOR POLAR */

	printf("\n");

	return 0;
}

// test kem fo
int test_kem_fo_cpucycles()
{
	printf("\n=============================test_kem_cpucycles==================================\n");
	unsigned long long t[NTESTS];
	unsigned char pk[CRYPTO_PUBLICKEYBYTES];
	unsigned char sk[CRYPTO_SECRETKEYBYTES];
	unsigned char k1[CRYPTO_SESSION], k2[CRYPTO_SESSION], c[CRYPTO_CIPHERTEXTBYTES];
	int i;

	for (i = 0; i < NTESTS; i++)
	{
		t[i] = cpucycles();
		crypto_kem_keypair(pk, sk);
	}
	printf("kem_fo kg  ");
	print_cpucycles(t);

	crypto_kem_keypair(pk, sk);
	random_bytes(k1, CRYPTO_SESSION);
	for (i = 0; i < NTESTS; i++)
	{
		t[i] = cpucycles();
		crypto_kem_enc(c, k1, pk);
	}
	printf("encryption ");
	print_cpucycles(t);

	/* FOR POLAR */
	// test polar_enc
	polar_enc_cpucycles();
	/* END FOR POLAR */

	crypto_kem_keypair(pk, sk);
	crypto_kem_enc(c, k1, pk);
	for (i = 0; i < NTESTS; i++)
	{
		t[i] = cpucycles();
		crypto_kem_dec(k2, c, sk);
	}
	printf("decryption ");
	print_cpucycles(t);

	/* FOR POLAR */
	// test polar_dec
	polar_dec_cpucycles();
	/* END FOR POLAR */

	printf("\n");

	return 0;
}

// test hash
int test_hash_cpucycles()
{
	int i, loop = 1000000;
	unsigned char buf[DIM_N], seed[SEED_LEN], out[MESSAGE_LEN];
	uint64_t t0, t1;
	random_bytes(seed, SEED_LEN);
	keccak_state state;
	shake256_absorb_once(&state, seed, SEED_LEN);
	shake256_squeeze(buf, DIM_N, &state);
	t0 = __rdtsc();
	for (i = 0; i < loop; i++)
	{
		// hash(buf, DIM_N, out);
		sha3_256(out, buf, DIM_N);
	}
	t1 = __rdtsc();
	printf("test hash speed:\n");
	//	printf("time: %f us\n",(t1-t0)/(2800.0*loop));
	printf("cpucycles: %f \n", (t1 - t0) / ((double)loop));
	printf("CPB: %f \n", (t1 - t0) / ((double)loop * DIM_N));
	printf("\n");

	return 0;
}


// test gen_psi
int test_gen_psi_cpucycles()
{
	int i, loop = 1000000;
	unsigned char seed[SEED_LEN];
	unsigned char e[DIM_N];
	uint64_t t0, t1;
	random_bytes(seed, SEED_LEN);
	keccak_state state;
    shake256_absorb_once(&state, seed, SEED_LEN);

	t0 = __rdtsc();
	for (i = 0; i < loop; i++)
	{
		gen_e(e, &state);
	}
	t1 = __rdtsc();
	printf("gen_e cpucycles: %f \n", (t1 - t0) / ((double)loop));
	printf("\n");

	return 0;
}

// test gen_a
int test_gen_a_cpucycles()
{
	int i, loop = 1000000;
	unsigned char seed[SEED_LEN];
	unsigned char a[DIM_N];
	uint64_t t0, t1, sum;
	random_bytes(seed, SEED_LEN);
	sum = 0;
	for (i = 0; i < loop; i++)
	{
		t0 = __rdtsc();
		gen_a(a, seed);
		t1 = __rdtsc();
		sum += (t1 - t0);
	}

	printf("test gen_a speed:  ");
	//	printf("time: %f us\n",(sum)/(2800.0*loop));
	printf("cpucycles: %f \n", (sum) / ((double)loop));
	printf("\n");

	return 0;
}

// test polymul
int test_poly_mul_cpucycles()
{
	int i, loop = 100000;
	unsigned char a[DIM_N], pk[DIM_N], seed[SEED_LEN];
	unsigned char sk[DIM_N];
	uint64_t t0, t1, sum;

	random_bytes(a, DIM_N);
	random_bytes(seed, SEED_LEN);
	keccak_state state;
    shake256_absorb_once(&state, seed, SEED_LEN);
	gen_e(sk, &state);

	sum = 0;
	for (i = 0; i < loop; i++)
	{
		t0 = __rdtsc();
		poly_mul(a, sk, pk, DIM_N);
		t1 = __rdtsc();
		sum += (t1 - t0);
	}

	printf("test poly_mul speed:\n");
	printf("cpucycles: %f \n", (sum) / ((double)loop));
	printf("\n");

	return 0;
}

int test_init()
{
	int i, loop = 10000;
	unsigned char a[DIM_N], pk[DIM_N], seed[SEED_LEN];
	unsigned char sk[DIM_N];

	random_bytes(a, DIM_N);
	random_bytes(seed, SEED_LEN);
	keccak_state state;
    shake256_absorb_once(&state, seed, SEED_LEN);
	gen_e(sk, &state);

	for (i = 0; i < loop; i++)
	{
		poly_mul(a, sk, pk, DIM_N);
	}

	return 0;
}
